import gradio as gr
import requests
import io
from util.utils import *
from results import results
from pathlib import Path
from step import Step

AccessKey, SecretKey = get_AccessKey_SecretKey()
site = 'https://openapi.liblibai.cloud/api/generate/comfyui/app'


class Step2(Step):
    PARAMS = {
        "templateUuid": "4df2efa0f18d46dc9758803e478eb51c",
        "generateParams": {
            "261": {
                "class_type": "LoadImage",
                "inputs": {
                    "image": None
                }
            },
            "269": {
                "class_type": "LoadImage",
                "inputs": {
                    "image": None
                }
            },
            "296": {
                "class_type": "ImpactInt",
                "inputs": {
                    "value": 1536
                }
            },
            "349": {
                "class_type": "RepeatLatentBatch",
                "inputs": {
                    "amount": 1
                }
            },
            "workflowUuid": "bbc223066f7649e9a2f6f4cf1f4bdb3c"
        }
    }

    def __init__(self):
        super().__init__("step2")
        

    def liblib(self,
        init_image: Image.Image,
        pose_image: Image.Image,
    ):
        if init_image is None:
            return None, "请先上传一张图片。"
        #upload images
        guesture_image_url = upload_image(AccessKey,SecretKey,init_image)
        background_image_url = upload_image(AccessKey,SecretKey,pose_image)
        print("init_image_url:", guesture_image_url)
        print("pose_image_url:", background_image_url)
        if guesture_image_url is None or background_image_url is None:
            return None, "图片上传失败，请重试。"

        m_params = self.PARAMS.copy()
        m_params["generateParams"]["261"]["inputs"]["image"] = guesture_image_url
        m_params["generateParams"]["269"]["inputs"]["image"] = background_image_url
        header = {
            'content-type': 'application/json'
        }
        sign, timestamp, signature_nonce = make_sign(SecretKey,type="generate")
        url = f"{site}?AccessKey={AccessKey}&Signature={sign}&SignatureNonce={signature_nonce}&Timestamp={timestamp}"
        response = requests.post(url, json=m_params, headers=header)
        if response.status_code == 200:
            result = response.json()
            print(result)
            if result.get("code") == 0:
                #get generate result
                generateUuid = result.get("data").get("generateUuid")
                sign, timestamp, signature_nonce = make_sign(SecretKey,type="status")
                get_url = "https://openapi.liblibai.cloud/api/generate/comfy/status"
                get_url = f"{get_url}?AccessKey={AccessKey}&Signature={sign}&SignatureNonce={signature_nonce}&Timestamp={timestamp}"
                data = {"generateUuid": generateUuid}
                while True:
                    resp = requests.post(get_url, json=data, headers=header)
                    if resp.status_code == 200:
                        res = resp.json()
                        if res.get("code") != 0:
                            continue
                        if res.get("data").get("generateStatus") == 5:
                            out = res.get("data").get("images")[0]['imageUrl']
                            print("generated img url:", out)
                            img_response = requests.get(out)
                            results.step2_result_image = Image.open(io.BytesIO(img_response.content))
                            self.save_img(results.step2_result_image, generateUuid)
                            return Image.open(io.BytesIO(img_response.content))
                        time.sleep(1)

        return None


    def create_step(self):
        with gr.Blocks(theme=gr.themes.Soft()) as demo:
            gr.Markdown("更换背景图像")

            with gr.Row():
                character_image = gr.Image(type="pil", label="上传人物图片", sources=["upload", "clipboard"], height=420)
                background_image = gr.Image(type="pil", label="上传背景图片", sources=["upload", "clipboard"], height=420)
                result_image = gr.Image(label="输出图像", interactive=False, height=420)


            run = gr.Button("生成", variant="primary")
            load = gr.Button("加载上一步图片",variant="primary")
            msg = gr.Markdown("")

            def _run(init_image, pose_image):
                global AccessKey,SecretKey
                AccessKey, SecretKey = get_AccessKey_SecretKey()
                if AccessKey is None or SecretKey is None:
                    gr.Warning("请先在[设置]中创建并选择一个用户以存储密钥",visible=True,duration=5)
                    return None
                out= self.liblib(init_image, pose_image)
                return out

            run.click(
                _run,
                inputs=[character_image, background_image],
                outputs=[result_image]
            )

            def _load():
                return results.step1_result_image

            load.click(
                _load,
                # inputs=[character_image],
                outputs=[character_image]
            )

        return demo

if __name__ == "__main__":
    step2 = Step2()
    step2.create_step().queue().launch()
